# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import json
import datetime
import pandas as pd
import numpy as np


class AutoInc(object):
  def __init__(self, start=1):
    self._idx = start - 1

  def __call__(self):
    self._idx += 1
    return self._idx


def load_feed_as_json(root_dir, trading_date, product):
  if not isinstance(trading_date, str):
    trading_date = trading_date.strftime('%Y%m%d')

  feed_file = root_dir.joinpath(f'{trading_date}/{product.full_symbol}')
  feed_list = []
  with feed_file.open() as in_file:
    for line in in_file:
      if line.strip():
        data = json.loads(line)
        feed_list.append(data)
  return feed_list


def load_multiday_feed_as_json(root_dir, trading_dates, product):
  res = []
  for trading_date in trading_dates:
    feed_list = load_feed_as_json(root_dir, trading_date, product)
    res.extend(feed_list)
  return res


def load_products_feed_as_json(root_dir, trading_dates, products):
  res = {}
  for product in products:
    res[product] = load_multiday_feed_as_json(root_dir, trading_dates, product)
  return res


def load_og_log(filename):
  data = json.load(open(filename))
  return data


def get_submit_price(og_event_list):
  buy_ts_list = []
  buy_price_list = []

  sell_ts_list = []
  sell_price_list = []
  for data in og_event_list:
    if data['type'] != 'ORDER_EVENT':
      continue
    event = data['event']

    if event['type'] == 'ORDER_SUBMITTED':
      if event['order_side'] == 'BUY_ORDER':
        buy_ts_list.append(int(event['event_time']))
        buy_price_list.append(event['order_price'])
      elif event['order_side'] == 'SELL_ORDER':
        sell_ts_list.append(int(event['event_time']))
        sell_price_list.append(event['order_price'])

  submit_buy_price = pd.Series(buy_price_list, index=pd.to_datetime(buy_ts_list))
  submit_sell_price = pd.Series(sell_price_list, index=pd.to_datetime(sell_ts_list))
  return submit_buy_price, submit_sell_price


def get_fill_price(og_event_list):
  buy_ts_list = []
  buy_price_list = []

  sell_ts_list = []
  sell_price_list = []
  for data in og_event_list:
    if data['type'] != 'ORDER_EVENT':
      continue
    event = data['event']

    if event['type'] == 'ORDER_FILLED':
      if event['order_side'] == 'BUY_ORDER':
        buy_ts_list.append(int(event['event_time']))
        buy_price_list.append(event['fill_price'])
      elif event['order_side'] == 'SELL_ORDER':
        sell_ts_list.append(int(event['event_time']))
        sell_price_list.append(event['fill_price'])

  fill_buy_price = pd.Series(buy_price_list, index=pd.to_datetime(buy_ts_list))
  fill_sell_price = pd.Series(sell_price_list, index=pd.to_datetime(sell_ts_list))
  return fill_buy_price, fill_sell_price


def get_fill_qty(og_event_list):
  buy_ts_list = []
  buy_qty_list = []

  sell_ts_list = []
  sell_qty_list = []
  for data in og_event_list:
    if data['type'] != 'ORDER_EVENT':
      continue
    event = data['event']

    if event['type'] == 'ORDER_FILLED':
      if event['order_side'] == 'BUY_ORDER':
        buy_ts_list.append(int(event['event_time']))
        buy_qty_list.append(event['fill_qty'])
      elif event['order_side'] == 'SELL_ORDER':
        sell_ts_list.append(int(event['event_time']))
        sell_qty_list.append(event['fill_qty'])

  fill_buy_qty = pd.Series(buy_qty_list, index=pd.to_datetime(buy_ts_list))
  fill_sell_qty = pd.Series(sell_qty_list, index=pd.to_datetime(sell_ts_list))
  return fill_buy_qty, fill_sell_qty


def get_mid_series(feed_list):
  ts_list = []
  mid_list = []
  for feed in feed_list:
    if feed.get('type') == 'book':
      mid = feed['data']['mid']
      ts = feed['timestamp']
      if not ts_list:
        ts_list.append(to_mid_night(ts))
        mid_list.append(np.nan)
      if ts_list[-1] != ts:
        ts_list.append(ts)
        mid_list.append(mid)
  mid_price_series = pd.Series(mid_list, index=pd.to_datetime(ts_list))
  return mid_price_series


def get_depth(feed_list, num_depth=5):
  ts_list = []
  buy_list = []
  sell_list = []

  def calc_depth(data):
    if len(data) >= num_depth:
      m = np.array(data[:num_depth])
    else:
      m = np.array(data)
    return sum(m[:, 0] * m[:, 1])

  for feed in feed_list:
    if feed.get('type') == 'book':
      if not ts_list:
        ts_list.append(to_mid_night(feed['timestamp']))
        buy_list.append(np.nan)
        sell_list.append(np.nan)

      ts_list.append(feed['timestamp'])
      buy_list.append(calc_depth(feed['data']['bids']))
      sell_list.append(calc_depth(feed['data']['asks']))
  buy_depth_series = pd.Series(buy_list, index=pd.to_datetime(ts_list))
  sell_depth_series = pd.Series(sell_list, index=pd.to_datetime(ts_list))
  return buy_depth_series, sell_depth_series


def get_trade_qty_series(feed_list):
  ts_list = []
  trade_qty_list = []
  prev = None

  for feed in feed_list:
    feed_type = feed.get('type')
    qty = feed.get('qty') or feed['data'].get('qty')
    if feed_type == 'trade' or qty:
      side = feed.get('side') or feed['data']['side']
      price = feed.get('price') or feed['data']['price']
      if 'SELL' in side:
        qty = -qty

      if not ts_list:
        ts_list.append(to_mid_night(feed['timestamp']))
        trade_qty_list.append(np.nan)

      ts = feed['timestamp']
      if ts == ts_list[-1]:
        trade_qty_list[-1] += qty
      else:
        ts_list.append(ts)
        trade_qty_list.append(qty)

  trade_qty_series = pd.Series(trade_qty_list, index=pd.to_datetime(ts_list))
  return trade_qty_series


def get_moving_average(time_series, window, resample=True):
  ser = time_series.rolling(window).mean()
  freq = datetime.timedelta(seconds=1)
  if resample:
    ser = ser.asfreq(freq, method='backfill')
  return ser


def get_moving_sum(time_series, window, resample=True):
  ser = time_series.rolling(window).sum()
  freq = datetime.timedelta(seconds=1)
  if resample:
    ser = ser.asfreq(freq, method='backfill')
  return ser


def get_moving_abs_sum(time_series, window, resample=True):
  ser = time_series.abs().rolling(window).sum()
  freq = datetime.timedelta(seconds=1)
  if resample:
    ser = ser.asfreq(freq, method='backfill')
  return ser


def get_moving_diff(time_series, window):
  return time_series.rolling(window).apply(lambda arr: arr[-1] - arr[0])


def get_moving_return(time_series, window):
  return time_series.rolling(window).apply(lambda arr: (arr[-1] - arr[0]) / arr[-1])


def to_mid_night(timestamp):
  x = pd.to_datetime(datetime.datetime.fromtimestamp(timestamp / 1e9).date())
  return int(x.timestamp() * 1e9)


def remove_nan_entries(ser):
  return ser[~pd.isna(ser)]
